# -*- coding: utf-8 -*-

from pg_connector import PostgresConnector
import argparse
import json


def main():
    parser = argparse.ArgumentParser(description='PostgreSQL数据库连接工具')
    parser.add_argument('--config', type=str, default='config.yml', help='配置文件路径')
    parser.add_argument('--action', type=str, required=True, 
                        choices=['list-dbs', 'list-tables', 'list-columns', 'list-all-columns', 
                                 'table-content', 'table-first-last'],
                        help='要执行的操作')
    parser.add_argument('--table', type=str, help='表名(用于查询列或内容)')
    parser.add_argument('--limit', type=int, default=10, help='查询限制条数')
    
    args = parser.parse_args()
    
    with PostgresConnector(args.config) as db:
        result = None
        
        if args.action == 'list-dbs':
            result = db.list_databases()
            print(f"数据库列表: {', '.join(result)}")
            
        elif args.action == 'list-tables':
            result = db.list_tables()
            print(f"表列表: {', '.join(result)}")
            
        elif args.action == 'list-columns':
            if not args.table:
                print("错误: 需要指定表名")
                return
            result = db.list_columns(args.table)
            print(f"表 '{args.table}' 的列:")
            for col in result:
                print(f"  {col['column']} ({col['type']})")
                
        elif args.action == 'list-all-columns':
            result = db.list_all_tables_columns()
            print("所有表的列:")
            for table, columns in result.items():
                print(f"表: {table}")
                for col in columns:
                    print(f"  {col['column']} ({col['type']})")
                print()
                
        elif args.action == 'table-content':
            if not args.table:
                print("错误: 需要指定表名")
                return
            result = db.get_table_content(args.table, args.limit)
            print(f"表 '{args.table}' 的内容 (限制 {args.limit} 条):")
            print(json.dumps(result, indent=2, ensure_ascii=False))
            
        elif args.action == 'table-first-last':
            if not args.table:
                print("错误: 需要指定表名")
                return
            result = db.get_table_first_last_records(args.table, args.limit)
            print(f"表 '{args.table}' 的前 {args.limit} 条记录:")
            print(json.dumps(result['first'], indent=2, ensure_ascii=False))
            print(f"\n表 '{args.table}' 的后 {args.limit} 条记录:")
            print(json.dumps(result['last'], indent=2, ensure_ascii=False))


if __name__ == "__main__":
    main() 